/*
 * Copyright (c) 2025, 大连理工大学 (Dalian University of Technology)
 *
 * Licensed under the Mulan PSL v2.
 * You can use this file according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of the License at
 *
 *     http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. SEE THE MULAN PSL v2
 * FOR MORE DETAILS.
 */

#pragma once

#include <unordered_map>
#include <utility>

#include "DBManager/Repository/HDF5/Base/SetRepositoryBase.h"

SG_DBMANAGER_HDF5_NAMESPACE_OPEN

template <typename T>

class DBManager_DECLSPEC SingleSetRepositoryBase : public SetRepositoryBase<T>
{
  public:
    explicit SingleSetRepositoryBase (HDF5Database& database) : SetRepositoryBase<T> (database) {};

    ~SingleSetRepositoryBase () {};

    std::vector<T>* LoadReadCache (std::size_t loadIndex) override
    {
        auto& mainCache       = this->GetMainCache ();
        auto& mainDatasetInfo = this->GetMainDatasetInfo ();

        std::size_t beginIndex = 0;
        std::size_t length     = 0;
        mainDatasetInfo.CalcIndexBeginAndLength (loadIndex, _OUT beginIndex, _OUT length, this->GetDatasetCacheSize ());

        auto* cacheVec = this->m_cacheHelper.template GetVectorFast<T> (mainCache.m_readCacheKey);
        if (!cacheVec)
        {
            cacheVec = new std::vector<T>;
            cacheVec->reserve (length);
            this->GetDataset (mainDatasetInfo.GetPath ()).select ({ beginIndex }, { length }).read (*cacheVec);
            this->m_cacheHelper.template CacheVector (mainCache.m_readCacheKey, cacheVec);
        }
        else
        {
            cacheVec->clear ();
            this->GetDataset (mainDatasetInfo.GetPath ()).select ({ beginIndex }, { length }).read (*cacheVec);
            // 更新该key在缓存池占用内存大小
            this->GetDatabase ().GetCachePool ().UpdateMemorySize (mainCache.m_readCacheKey, cacheVec->size () * sizeof (T));
        }
        mainCache.dataMinIndex = beginIndex;
        mainCache.dataMaxIndex = beginIndex + length;
        return cacheVec;
    }

    // 转存数据
    void DumpWriteCahce () override
    {
        // 如果缓存区没有数据，直接返回
        if (this->GetMainCache ().m_writeStorage.empty ())
        {
            return;
        }
        // 将主表的顺序标记设为 false，表示数据可能不是顺序存储的
        this->GetMainDatasetInfo ().SetIsSequential (false);

        // 将缓存区中的数据追加到主表中
        this->WriteData (this->GetMainDatasetInfo (), this->GetMainCache ().m_writeStorage);

        // 更新Index表
        this->WriteIndexData (this->GetIndexDatasetInfo (), this->GetMainCache ().m_writeStorage);

        // 清空缓存区的数据
        this->GetMainCache ().m_writeStorage.clear ();
    }

    /**
     * @brief 纯虚函数，用于对数据进行排序，需由派生类实现。
     *
     * @return 排序是否成功
     */
    void SortData () override
    {
        if (!this->GetMainCache ().m_writeStorage.empty ())
        {
            this->DumpWriteCahce ();
        }

        auto& mainInfo = this->GetMainDatasetInfo ();

        if (mainInfo.IsSequential ())
        {
            return;
        }

        auto mainDataset = this->GetDataset (mainInfo.GetPath ());

        std::vector<T> datas;

        mainDataset.read (datas);

        // 对数据进行排序
        std::sort (datas.begin (), datas.end ());

        mainDataset.select ({ 0 }, { datas.size () }).write (datas);

        this->RewriteIndexData (datas);

        // 设置为已顺序
        mainInfo.SetIsSequential (true);

        const std::string rangeCacheKey = mainInfo.GetPath () + std::string ("_SIdRangeCache");
        const std::string cacheKeyBase  = mainInfo.GetPath () + std::string ("_SIdRangeDataCache_");
        this->GetDatabase ().GetCachePool ().Remove (rangeCacheKey);
        this->GetDatabase ().GetCachePool ().Remove (cacheKeyBase);
    }

    /**
     * @brief   : 根据 SID 查找实体
     * @param   : sId
     * @param   : arrs
     * @return  :
     */
    SG::DataStructure::Common::Status FindBySId (SG::DataStructure::Common::Id sId, std::vector<T>& _OUT arrs) override
    {
        // 1. 确定数据是排序过的
        this->Sort ();

        if (!this->GetMainDatasetInfo ().IsCreate ())
        {
            return SG::DataStructure::Common::Status::STATUS_Failure;
        }

        // 2. 尝试从缓存中获取所有 SID 的范围映射：Id -> (startIndex, length)
        //    与 FindById 的缓存思路一致，但这里缓存的是每个 SID 所在的连续区间
        auto& mainInfo = this->GetMainDatasetInfo ();
        auto& idxInfo  = this->GetIndexDatasetInfo ();

        const std::string rangeCacheKey = mainInfo.GetPath () + std::string ("_SIdRangeCache");

        using RangePair  = std::pair<std::size_t, std::size_t>;
        auto* rangeCache = this->m_cacheHelper.template GetUnorderedMapFast<SG::DataStructure::Common::Id, RangePair> (rangeCacheKey);

        if (rangeCache == nullptr)
        {
            // 未命中缓存：读取所有索引ID，并构建所有 SID -> 区间 的映射
            std::vector<SG::DataStructure::Common::Id> idList;
            this->GetDataset (idxInfo).read (idList);

            // 先统计唯一 SID 的数量（idList 已排序）。使用严格弱序定义的等价判断，避免对 operator== 的依赖
            std::size_t uniqueCount = 0;
            {
                std::size_t j = 0;
                while (j < idList.size ())
                {
                    ++uniqueCount;
                    const auto& v = idList[j];
                    while (j < idList.size () && !(idList[j] < v) && !(v < idList[j]))
                    {
                        ++j;
                    }
                }
            }

            auto* newMap = new std::unordered_map<SG::DataStructure::Common::Id, RangePair> ();
            newMap->reserve (uniqueCount);

            std::size_t i = 0;
            while (i < idList.size ())
            {
                const auto&       curId = idList[i];
                const std::size_t st    = i;
                while (i < idList.size () && !(idList[i] < curId) && !(curId < idList[i]))
                {
                    ++i;
                }
                const std::size_t len = i - st;
                (*newMap)[curId]      = { st, len };
            }

            this->m_cacheHelper.template CacheUnorderedMap<SG::DataStructure::Common::Id, RangePair> (rangeCacheKey, newMap);
            rangeCache = newMap;
        }

        // 3. 根据 sId 查找范围
        auto it = rangeCache->find (sId);
        if (it == rangeCache->end ())
        {
            return SG::DataStructure::Common::Status::STATUS_Failure;
        }

        const std::size_t startIndex = it->second.first;
        const std::size_t length     = it->second.second;
        if (length == 0)
        {
            return SG::DataStructure::Common::Status::STATUS_Failure;
        }

        // 4. 优先尝试通过主读缓存获取数据；若缓存未覆盖则触发加载，必要时分多段从缓存拼接
        arrs.clear ();
        arrs.reserve (length);

        auto&             mainCache = this->GetMainCache ();
        std::vector<T>*   cacheVec  = this->m_cacheHelper.template GetVectorFast<T> (mainCache.m_readCacheKey);
        std::size_t       copied    = 0;
        const std::size_t endIndex  = startIndex + length;  // 需要的右开边界

        auto copyFromCache = [&] (std::size_t segBegin, std::size_t segEnd) {
            if (segBegin >= segEnd)
                return;
            if (!cacheVec)
                return;
            const std::size_t off   = segBegin - mainCache.dataMinIndex;
            const std::size_t count = segEnd - segBegin;
            arrs.insert (arrs.end (), cacheVec->begin () + off, cacheVec->begin () + off + count);
        };

        while (copied < length)
        {
            // 如果当前缓存不覆盖下一个需要的位置，则加载包含该位置的读缓存
            const std::size_t nextIndex = startIndex + copied;
            if (nextIndex < mainCache.dataMinIndex || nextIndex >= mainCache.dataMaxIndex
                || (mainCache.dataMinIndex == 0 && mainCache.dataMaxIndex == 0))
            {
                cacheVec = this->LoadReadCache (nextIndex);
                if (!cacheVec)
                {
                    // 缓存加载失败，回退到直接读取该剩余区间，但需追加到 arrs（而非覆盖）
                    const std::size_t remain = length - copied;
                    std::vector<T>    tmp;
                    tmp.reserve (remain);
                    this->GetDataset (mainInfo).select ({ nextIndex }, { remain }).read (tmp);
                    arrs.insert (arrs.end (), tmp.begin (), tmp.end ());
                    copied += remain;
                    break;
                }
            }

            // 计算当前缓存可覆盖的区间，与还需拷贝的区间的交集
            const std::size_t segBegin = std::max (nextIndex, mainCache.dataMinIndex);
            const std::size_t segEnd   = std::min (endIndex, mainCache.dataMaxIndex);
            if (segBegin >= segEnd)
            {
                // 理论上不应发生，防御性处理：强制加载从 nextIndex 开始的缓存
                cacheVec = this->LoadReadCache (nextIndex);
                if (!cacheVec)
                {
                    const std::size_t remain = length - copied;
                    std::vector<T>    tmp;
                    tmp.reserve (remain);
                    this->GetDataset (mainInfo).select ({ nextIndex }, { remain }).read (tmp);
                    arrs.insert (arrs.end (), tmp.begin (), tmp.end ());
                    copied += remain;
                    break;
                }
                continue;
            }

            // 从缓存复制该段
            copyFromCache (segBegin, segEnd);
            copied += (segEnd - segBegin);
        }

        return SG::DataStructure::Common::Status::STATUS_Success;
    }
};

SG_DBMANAGER_HDF5_NAMESPACE_CLOSE